# -*- coding: UTF-8 -*-
import torch
import config


class Generator(torch.nn.Module):
    def __init__(self, n_node, node_emd_init, manifold, c):
        super().__init__()
        self.n_node = n_node
        self.node_emd_init = node_emd_init
        self.manifold = manifold
        self.c = c

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        self.embedding_matrix = torch.nn.Parameter(
            torch.tensor(self.node_emd_init, device=self.device, requires_grad=True))
        self.bias_vector = torch.nn.Parameter(torch.zeros(self.n_node, device=self.device))

        # self.all_score = torch.matmul(self.embedding_matrix, self.embedding_matrix_transpose) + self.bias_vector
        self.all_score = self.forward()

        self.reward = None
        self.g_loss = None

        self.optimizer = torch.optim.Adam([self.embedding_matrix, self.bias_vector], lr=config.lr_gen)

        self.node_embedding = None
        self.node_neighbor_embedding = None
        self.bias = None
        self.score = None
        self.prob = None

        self.hyp_node_embedding = None
        self.hyp_node_neighbor_embedding = None
        self.hyp_bias = None
        self.hyp_score = None
        self.hyp_score_f = None

        self.hyp_embedding_matrix = None
        self.hyp_bias_vector = None
        self.hyp_all_score = None

    def train(self, node_id, node_neighbor_id, reward):
        self.reward = reward
        self.node_embedding = self.embedding_matrix[node_id]
        self.node_neighbor_embedding = self.embedding_matrix[node_neighbor_id]
        self.bias = self.bias_vector[node_neighbor_id]
        self.bias[torch.isnan(self.bias)] = 0

        # 将欧式特征向量转换为非欧向量
        self.hyp_node_embedding = torch.cat((torch.zeros(len(node_id), 1).to(self.device), self.node_embedding), dim=1)
        self.hyp_node_embedding = self.manifold.expmap0(self.hyp_node_embedding, self.c)
        self.hyp_node_neighbor_embedding = torch.cat(
            (torch.zeros(len(node_neighbor_id), 1).to(self.device), self.node_neighbor_embedding), dim=1)
        self.hyp_node_neighbor_embedding = self.manifold.expmap0(self.hyp_node_neighbor_embedding, self.c)

        # 偏移量增加维度
        self.hyp_bias = self.manifold.proj_tan0(self.bias.view(1, -1), self.c)

        # 进行非欧计算
        self.hyp_score = self.manifold.mobius_matvec(self.hyp_node_embedding, self.hyp_node_neighbor_embedding, self.c)
        self.hyp_score = self.manifold.proj(self.hyp_score, self.c)
        self.hyp_score = self.manifold.mobius_add(self.hyp_score, self.hyp_bias, self.c)
        self.hyp_score = self.manifold.proj(self.hyp_score, self.c)

        # 转换回欧式向量
        self.score = self.manifold.logmap0(self.hyp_score, self.c)
        self.score = torch.sum(self.score, 1)
        self.prob = torch.clamp(torch.sigmoid(self.score), 1e-5, 1)

        self.optimizer.zero_grad()
        l1_loss = -torch.mean(torch.log(self.prob) * self.reward)
        l2 = config.lambda_gen * (0.5 * torch.sum(self.node_neighbor_embedding ** 2) +
                                  0.5 * torch.sum(self.node_embedding ** 2) +
                                  0.5 * torch.sum(self.bias ** 2))
        self.g_loss = l1_loss + l2

        self.g_loss.backward()
        self.optimizer.step()

        return self.g_loss

    def forward(self):
        with torch.no_grad():
            # self.all_score = torch.matmul(self.embedding_matrix, self.embedding_matrix_transpose) + self.bias_vector
            self.hyp_embedding_matrix = torch.cat(
                (torch.zeros(len(self.embedding_matrix), 1).to(self.device), self.embedding_matrix), dim=1)
            self.hyp_embedding_matrix = self.manifold.expmap0(self.hyp_embedding_matrix, self.c)
            self.hyp_bias_vector = self.manifold.proj_tan0(self.bias_vector.view(1, -1), self.c)

            self.hyp_all_score = self.manifold.mobius_matvec(self.hyp_embedding_matrix, self.hyp_embedding_matrix,
                                                             self.c)
            self.hyp_all_score = self.manifold.proj(self.hyp_all_score, self.c)
            self.hyp_all_score = self.manifold.mobius_add(self.hyp_all_score, self.hyp_bias_vector, c=self.c)
            self.hyp_all_score = self.manifold.proj(self.hyp_all_score, self.c)

            self.all_score = self.manifold.logmap0(self.hyp_all_score, self.c)
            return self.all_score

    def get_all_score(self):
        return self.all_score
